import os
import re
import pandas as pd
import matplotlib.pyplot as plt

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

# 文件夹路径
DP_DIR = 'output/output/dp'
GREEDY_DIR = 'output/output/greedy'

# 需要对比的背包容量
CAPACITIES = [10000, 100000, 1000000]

# 匹配文件名的正则表达式
dp_pattern = re.compile(r'dp_(\d+)_(' + '|'.join(str(c) for c in CAPACITIES) + r')\.csv')
greedy_pattern = re.compile(r'greedy_(\d+)_(' + '|'.join(str(c) for c in CAPACITIES) + r')\.csv')

# 收集文件
def collect_files(folder, pattern):
    files = {}
    if not os.path.exists(folder):
        print(f"文件夹不存在: {folder}")
        return files
    for fname in os.listdir(folder):
        match = pattern.match(fname)
        if match:
            item_count = int(match.group(1))
            capacity = int(match.group(2))
            if capacity not in files:
                files[capacity] = {}
            files[capacity][item_count] = os.path.join(folder, fname)
    print(f"在 {folder} 中找到 {len(files)} 个容量组")
    for capacity, items in files.items():
        print(f"  容量 {capacity}: {len(items)} 个文件")
    return files

dp_files = collect_files(DP_DIR, dp_pattern)
greedy_files = collect_files(GREEDY_DIR, greedy_pattern)

def get_time_from_csv(path):
    # 读取最后一行，提取耗时（毫秒），不再转换为秒
    try:
        with open(path, 'r', encoding='utf-8') as f:
            lines = f.readlines()
            for line in reversed(lines):
                if line.startswith('耗时(ms)'):
                    ms = float(line.strip().split(',')[1])
                    return ms
        return None
    except Exception as e:
        print(f"读取文件 {path} 时出错: {e}")
        return None

# 绘图
for capacity in CAPACITIES:
    dp_data = []
    greedy_data = []
    # 物品数量排序
    item_counts = sorted(set(dp_files.get(capacity, {}).keys()) | set(greedy_files.get(capacity, {}).keys()))
    print(f"\n处理容量 {capacity}，找到 {len(item_counts)} 个物品数量")
    
    for item_count in item_counts:
        # DP
        dp_time = None
        if item_count in dp_files.get(capacity, {}):
            dp_time = get_time_from_csv(dp_files[capacity][item_count])
        # Greedy
        greedy_time = None
        if item_count in greedy_files.get(capacity, {}):
            greedy_time = get_time_from_csv(greedy_files[capacity][item_count])
        if dp_time is not None:
            dp_data.append((item_count, dp_time))
        if greedy_time is not None:
            greedy_data.append((item_count, greedy_time))
    
    print(f"容量 {capacity}: DP数据点 {len(dp_data)}, Greedy数据点 {len(greedy_data)}")
    
    # 排序
    dp_data.sort()
    greedy_data.sort()
    # 分别取x和y
    dp_x = [x for x, _ in dp_data]
    dp_y = [y for _, y in dp_data]
    greedy_x = [x for x, _ in greedy_data]
    greedy_y = [y for _, y in greedy_data]
    
    # 检查是否有数据
    if not dp_data and not greedy_data:
        print(f"警告: 容量 {capacity} 没有任何数据，跳过绘图")
        continue
    
    # 绘图
    plt.figure()
    if dp_data:
        plt.plot(dp_x, dp_y, marker='o', label='动态规划')
    if greedy_data:
        plt.plot(greedy_x, greedy_y, marker='s', label='贪心算法')
    plt.xlabel('物品数量')
    plt.ylabel('运行时间（毫秒）')
    plt.title(f'背包容量={capacity}时算法运行时间对比')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(f'output/compare_{capacity}.png')
    plt.show()

# =================== 蛮力法与回溯法 ===================
BRUTE_DIR = 'output/output/brute_force'
BACKTRACK_DIR = 'output/output/backtrack'
CAPACITIES_BB = [100, 300, 500]
brute_pattern = re.compile(r'brute_force_(\d+)_(' + '|'.join(str(c) for c in CAPACITIES_BB) + r')\.csv')
backtrack_pattern = re.compile(r'backtrack_(\d+)_(' + '|'.join(str(c) for c in CAPACITIES_BB) + r')\.csv')

def collect_files_bb(folder, pattern):
    files = {}
    if not os.path.exists(folder):
        return files
    for fname in os.listdir(folder):
        match = pattern.match(fname)
        if match:
            item_count = int(match.group(1))
            capacity = int(match.group(2))
            if capacity not in files:
                files[capacity] = {}
            files[capacity][item_count] = os.path.join(folder, fname)
    return files

brute_files = collect_files_bb(BRUTE_DIR, brute_pattern)
backtrack_files = collect_files_bb(BACKTRACK_DIR, backtrack_pattern)

for capacity in CAPACITIES_BB:
    brute_data = []
    backtrack_data = []
    item_counts = sorted(set(brute_files.get(capacity, {}).keys()) | set(backtrack_files.get(capacity, {}).keys()))
    for item_count in item_counts:
        brute_time = None
        if item_count in brute_files.get(capacity, {}):
            brute_time = get_time_from_csv(brute_files[capacity][item_count])
        backtrack_time = None
        if item_count in backtrack_files.get(capacity, {}):
            backtrack_time = get_time_from_csv(backtrack_files[capacity][item_count])
        if brute_time is not None:
            brute_data.append((item_count, brute_time))
        if backtrack_time is not None:
            backtrack_data.append((item_count, backtrack_time))
    # 排序（按物品数量升序）
    brute_data.sort()
    backtrack_data.sort()
    # 分别取x和y
    brute_x = [x for x, y in brute_data]
    brute_y = [y for x, y in brute_data]
    backtrack_x = [x for x, y in backtrack_data]
    backtrack_y = [y for x, y in backtrack_data]
    # 绘图
    plt.figure()
    plt.plot(brute_x, brute_y, marker='o', label='蛮力法')
    plt.plot(backtrack_x, backtrack_y, marker='s', label='回溯法')
    plt.xlabel('物品数量')
    plt.ylabel('运行时间（毫秒）')
    plt.title(f'背包容量={capacity}时蛮力法与回溯法运行时间对比')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(f'output/compare_bb_{capacity}.png')
    plt.show()
